# prophet_multivar_4day.py
import argparse, warnings, re
warnings.filterwarnings("ignore")
import pandas as pd
import numpy as np
from pandas.tseries.offsets import MonthEnd

TRAIN_YEARS = list(range(2017, 2023))
FORE_YEARS  = [2025, 2026]

def make_4day_grid_for_year(year: int):
    d = pd.Timestamp(year=year, month=3, day=1)
    end = pd.Timestamp(year=year, month=7, day=1)
    out = []
    while d <= end:
        out.append(d); d += pd.Timedelta(days=4)
    return pd.DatetimeIndex(out)

def build_future_grid():
    return pd.DatetimeIndex(sorted(np.concatenate([make_4day_grid_for_year(y) for y in FORE_YEARS])))

def ensure_percent_scale(s: pd.Series) -> pd.Series:
    s = pd.to_numeric(s, errors="coerce")
    if s.max(skipna=True) <= 1.0: s = s * 100.0
    return s.clip(lower=0, upper=100)

# ---------- 训练表 ----------
def load_train_table(path, target_col, cols=None):
    df = pd.read_csv(path)
    # 时间
    date_col = None
    for c in ["date","时间","日期","ds","time"]:
        if c in df.columns: date_col = c; break
    if date_col is None:
        raise ValueError("训练表未找到时间列（date/时间/日期/ds/time）")
    df["date"] = pd.to_datetime(df[date_col])
    # 仅 2017–2022
    df = df[(df["date"].dt.year>=min(TRAIN_YEARS)) & (df["date"].dt.year<=max(TRAIN_YEARS))].copy()
    # 目标
    if target_col not in df.columns:
        raise ValueError(f"找不到目标列 {target_col}")
    df[target_col] = ensure_percent_scale(df[target_col])
    # 自变量列
    if cols is None:
        num_cols = df.select_dtypes(include=[np.number]).columns.tolist()
        cols = [c for c in num_cols if c != target_col]
    # 组装
    train = df[["date", target_col] + cols].copy().rename(columns={"date":"ds", target_col:"y"})
    return train, cols

# ---------- 解析 all_vars（季度/年月/日期 皆可） ----------
def _try_parse_quarter_str(s: pd.Series) -> pd.Series:
    def _to_dt(x):
        if pd.isna(x): return pd.NaT
        xs = str(x).strip()
        zh_map = {"一":"1","二":"2","三":"3","四":"4","Ⅰ":"1","Ⅱ":"2","Ⅲ":"3","Ⅳ":"4"}
        for k,v in zh_map.items(): xs = xs.replace(k, v)
        m = re.match(r"^\s*(\d{4})\s*[-/ ]?\s*[Q第]?\s*([1-4])\s*[季度Q]?\s*$", xs)
        if m:
            y = int(m.group(1)); q = int(m.group(2))
            return pd.Timestamp(year=y, month=q*3, day=1) + MonthEnd(0)
        return pd.NaT
    return s.map(_to_dt)

def _try_parse_yearmonth(s: pd.Series) -> pd.Series:
    def _to_dt(x):
        if pd.isna(x): return pd.NaT
        xs = str(x).strip()
        if re.match(r"^\d{6}$", xs):
            dt = pd.to_datetime(xs, format="%Y%m", errors="coerce")
            return dt + MonthEnd(0) if pd.notna(dt) else pd.NaT
        if re.match(r"^\d{4}[-/.]\d{1,2}$", xs):
            dt = pd.to_datetime(xs, errors="coerce")
            return (dt + MonthEnd(0)) if pd.notna(dt) else pd.NaT
        return pd.NaT
    return s.map(_to_dt)

def _coerce_datetime(df: pd.DataFrame, date_col: str) -> pd.DatetimeIndex:
    ser = df[date_col]
    dt = pd.to_datetime(ser, errors="coerce", infer_datetime_format=True)
    if dt.notna().mean() > 0.8: return pd.DatetimeIndex(dt)
    dt = _try_parse_quarter_str(ser)
    if dt.notna().mean() > 0.8: return pd.DatetimeIndex(dt)
    dt = _try_parse_yearmonth(ser)
    if dt.notna().mean() > 0.8: return pd.DatetimeIndex(dt)
    return pd.DatetimeIndex([pd.NaT]*len(ser))

def guess_and_parse_datetime(df: pd.DataFrame) -> pd.DatetimeIndex:
    for cand in ["date","时间","日期","ds","time"]:
        if cand in df.columns:
            dt = _coerce_datetime(df, cand)
            if dt.notna().any(): return dt
    low = {c.lower(): c for c in df.columns}
    if "year" in low and ("quarter" in low or "season" in low):
        y = pd.to_numeric(df[low["year"]], errors="coerce")
        qn = low["quarter"] if "quarter" in low else low["season"]
        q = pd.to_numeric(df[qn], errors="coerce").clip(1,4)
        dt = pd.to_datetime(dict(year=y, month=(q*3).astype("Int64"), day=1), errors="coerce") + MonthEnd(0)
        return pd.DatetimeIndex(dt)
    if "year" in low and "month" in low:
        y = pd.to_numeric(df[low["year"]], errors="coerce")
        m = pd.to_numeric(df[low["month"]], errors="coerce").clip(1,12)
        dt = pd.to_datetime(dict(year=y, month=m, day=1), errors="coerce") + MonthEnd(0)
        return pd.DatetimeIndex(dt)
    first = df.columns[0]
    dt = _coerce_datetime(df, first)
    if dt.notna().any(): return dt
    raise ValueError("无法识别时间列；可用列名：" + ", ".join(map(str, df.columns)))

# ---------- 因子外推（2025–2026 的 4天栅格） ----------
def forecast_quarterly_const(series_quarter: pd.Series, fut_grid: pd.DatetimeIndex) -> pd.Series:
    s = series_quarter.copy().sort_index().resample("Q").mean().dropna()
    try:
        from prophet import Prophet
        dfp = pd.DataFrame({"ds": s.index, "y": s.values})
        m = Prophet(weekly_seasonality=False, daily_seasonality=False, yearly_seasonality=True)
        m.fit(dfp)
        q_targets = []
        for y in FORE_YEARS:
            q_targets += [pd.Timestamp(f"{y}-03-31"), pd.Timestamp(f"{y}-06-30")]
        futq = pd.DataFrame({"ds": pd.to_datetime(q_targets)})
        yhat = m.predict(futq)[["ds","yhat"]].set_index("ds")["yhat"]
    except Exception:
        idx = s.index.view("i8").astype(float)
        vals = s.values.astype(float)
        k, b = np.polyfit(idx, vals, 1)
        q_targets = []
        for y in FORE_YEARS:
            q_targets += [pd.Timestamp(f"{y}-03-31"), pd.Timestamp(f"{y}-06-30")]
        idx2 = pd.to_datetime(q_targets).view("i8").astype(float)
        yhat = pd.Series(k*idx2 + b, index=pd.to_datetime(q_targets))

    def quarter_end(ts):
        q = (ts.month-1)//3 + 1
        return pd.Timestamp(year=ts.year, month=q*3, day=1) + MonthEnd(0)
    qe = pd.Series([quarter_end(d) for d in fut_grid], index=fut_grid)
    yhat_full = yhat.reindex(yhat.index.union(qe.unique())).sort_index().ffill().bfill()
    return yhat_full.reindex(qe.values).set_axis(fut_grid)

def forecast_continuous(series_daily: pd.Series, fut_grid: pd.DatetimeIndex) -> pd.Series:
    s = pd.to_numeric(series_daily, errors="coerce").asfreq("D").interpolate("time").ffill().bfill()
    s_train = s[(s.index.year>=min(TRAIN_YEARS)) & (s.index.year<=max(TRAIN_YEARS))]
    try:
        from prophet import Prophet
        dfp = pd.DataFrame({"ds": s_train.index, "y": s_train.values})
        m = Prophet(weekly_seasonality=False, daily_seasonality=False, yearly_seasonality=True)
        m.fit(dfp)
        fut = pd.DataFrame({"ds": fut_grid})
        pred = m.predict(fut)[["ds","yhat"]].set_index("ds")["yhat"]
        return pred
    except Exception:
        # 退化：线性 + 年周期
        t0 = (s_train.index - s_train.index[0]).days.values.astype(float)
        y  = s_train.values.astype(float)
        w = 2*np.pi/365.25
        X = np.c_[np.ones_like(t0), t0, np.sin(w*t0), np.cos(w*t0)]
        beta = np.linalg.pinv(X) @ y
        t_pred = (fut_grid - s_train.index[0]).days.values.astype(float)
        Xp = np.c_[np.ones_like(t_pred), t_pred, np.sin(w*t_pred), np.cos(w*t_pred)]
        yhat = Xp @ beta
        return pd.Series(yhat, index=fut_grid, dtype=float)

def build_future_regressors(allvars_csv, cols, grad_col, fut_grid):
    raw = pd.read_csv(allvars_csv)
    dt = guess_and_parse_datetime(raw)
    raw = raw.copy()
    raw["date"] = dt
    raw = raw.dropna(subset=["date"]).sort_values("date").set_index("date")
    num = raw.select_dtypes(include=[np.number]).copy()
    if num.empty:
        raise ValueError("all_vars_yoy.csv 中没有数值列。")
    daily_all = num.resample("D").mean().interpolate("time").ffill().bfill()

    fac = pd.DataFrame(index=fut_grid)
    for col in cols:
        if col not in daily_all.columns:
            raise ValueError(f"列 `{col}` 不在 all_vars_yoy.csv 的数值列中。")
        if grad_col and col == grad_col:
            fac[col] = forecast_quarterly_const(num[col], fut_grid)
        else:
            fac[col] = forecast_continuous(daily_all[col], fut_grid)
    return fac.reset_index().rename(columns={"index":"date"})

# ---------- Prophet 多元 ----------
def prophet_with_regressors(train_df: pd.DataFrame, future_df: pd.DataFrame) -> pd.DataFrame:
    from prophet import Prophet
    reg_cols = [c for c in train_df.columns if c not in {"ds","y"}]
    # 清洗缺失
    before = len(train_df)
    train_df = train_df.dropna(subset=["y"] + reg_cols)
    if before - len(train_df) > 0:
        print(f"[INFO] dropped {before - len(train_df)} NaN rows for training")

    m = Prophet(weekly_seasonality=False, daily_seasonality=False, yearly_seasonality=True)
    for c in reg_cols:
        m.add_regressor(c, standardize=True)
    m.fit(train_df[["ds","y"]+reg_cols])

    # 未来缺失兜底
    future_df = future_df.copy()
    for c in reg_cols:
        if future_df[c].isna().any():
            future_df[c] = future_df[c].interpolate().ffill().bfill()

    yhat = m.predict(future_df[["ds"]+reg_cols])[["ds","yhat"]]
    yhat["yhat"] = yhat["yhat"].clip(0,100)
    return yhat.rename(columns={"ds":"date","yhat":"就业率（%）"})

def main():
    ap = argparse.ArgumentParser(description="Prophet 多元（3/1~7/1 每4天；2017-2022 训练，2025-2026 预测）")
    ap.add_argument("--train", required=True, help="训练表 CSV（建议用 pre_2016_2022_grid.csv）")
    ap.add_argument("--allvars", required=True, help="外生变量全量 CSV（用于因子外推）")
    ap.add_argument("--target-col", default="employment_rate")
    ap.add_argument("--grad-col", required=True, help="毕业生人数列名（季度常数，不插值）")
    ap.add_argument("--cols", nargs="*", default=None, help="参与建模/外推的因子列（至少4个）")
    ap.add_argument("--out", default="prophet_forecast_2025_2026_4day.csv")
    args = ap.parse_args()

    train, cols = load_train_table(args.train, args.target_col, args.cols)
    fut_idx = build_future_grid()
    factors_future = build_future_regressors(args.allvars, cols, args.grad_col, fut_idx)

    # 训练与预测
    train_df = train.rename(columns={"date":"ds"}) if "date" in train.columns else train
    fut_df = factors_future.rename(columns={"date":"ds"})
    yhat = prophet_with_regressors(train_df, fut_df)
    yhat.to_csv(args.out, index=False, encoding="utf-8-sig")
    print(f"[Prophet] saved: {args.out} rows={len(yhat)}")

if __name__ == "__main__":
    main()
